package com.example.springbootdemo.text.en;

import org.apache.commons.lang3.StringUtils;
import org.tartarus.snowball.ext.EnglishStemmer;

import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

public class SplitWordUtil {
    public static Set<String> split(String sentence) {
        // 定义停用词列表
        Set<String> stopWords = new HashSet<>(Arrays.asList("a", "an", "the", "is", "at", "which", "on", "of"));

        // 移除标点符号
        String cleanedSentence = StringUtils.stripAccents(sentence.replaceAll("[,.\\-]", " ")).toLowerCase();

        // 分词
        List<String> words = Arrays.asList(cleanedSentence.split("\\s+"));

        // 词干化
        Set<String> stemmedWords = stemWords(words);

        // 移除停用词
        Set<String> filteredWords = removeStopWords(stemmedWords, stopWords);

        return filteredWords;
    }

    private static Set<String> stemWords(List<String> words) {
        Set<String> stemmedWords = new HashSet<>();
        EnglishStemmer stemmer = new EnglishStemmer();

        for (String word : words) {
            stemmer.setCurrent(word);
            if (stemmer.stem()) {
                stemmedWords.add(stemmer.getCurrent());
            } else {
                stemmedWords.add(word);
            }
        }

        return stemmedWords;
    }

    private static Set<String> removeStopWords(Set<String> words, Set<String> stopWords) {
        Set<String> filteredWords = new HashSet<>();
        for (String word : words) {
            if (!stopWords.contains(word)) {
                filteredWords.add(word);
            }
        }
        return filteredWords;
    }
}
